import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import models
from torch.utils.data import DataLoader, TensorDataset
import numpy as np
# 导入必要的库
'''
ResNet（Residual Network）通常用于图像分类任务，因为它能够处理复杂的高层特征表示。
然而，如果你希望使用ResNet来训练一个非线性回归任务，你可以通过一些修改来实现这一目标。
以下是一个基本的步骤指南，展示如何使用PyTorch和ResNet进行非线性回归任务。
'''

# 准备数据
'''
假设你已经有了特征和目标数据，分别存储在X_train和y_train中。你需要将它们转换为PyTorch张量，并创建数据加载器。
'''
# 假设 X_train 和 y_train 是 numpy 数组
X_train = np.random.rand(100, 3, 224, 224).astype(np.float32)  # 示例数据：100 张 3x224x224 的图像
y_train = np.random.rand(100, 1).astype(np.float32)             # 示例目标：100 个一维标签

# 转换为 PyTorch 张量
X_train_tensor = torch.tensor(X_train)
y_train_tensor = torch.tensor(y_train)

# 创建数据加载器
train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

# 修改ResNet模型
# 你需要修改ResNet模型的最后几层，以适应回归任务（通常是去除分类层并添加一个新的线性层）。
# 加载预训练的 ResNet 模型
resnet = models.resnet18(pretrained=True)

# 冻结所有参数（可选，如果你想微调，可以解冻部分层）
for param in resnet.parameters():
    param.requires_grad = False

# 修改最后的全连接层
num_ftrs = resnet.fc.in_features
resnet.fc = nn.Linear(num_ftrs, 1)  # 回归任务，输出一个值

# 定义损失函数和优化器
criterion = nn.MSELoss()  # 均方误差损失函数，适合回归任务
optimizer = optim.Adam(resnet.fc.parameters(), lr=0.001)  # 只优化最后的全连接层

# 训练模型
num_epochs = 20

for epoch in range(num_epochs):
    running_loss = 0.0
    for inputs, targets in train_loader:
        optimizer.zero_grad()

        # 前向传播
        outputs = resnet(inputs)
        loss = criterion(outputs, targets)

        # 反向传播和优化
        loss.backward()
        optimizer.step()

        running_loss += loss.item() * inputs.size(0)

    epoch_loss = running_loss / len(train_loader.dataset)
    print(f'Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.4f}')

#  测试模型（可选）
# 如果你有测试集，你可以按照类似的方式评估模型的性能。
# 假设 X_test 和 y_test 是 numpy 数组
X_test = np.random.rand(20, 3, 224, 224).astype(np.float32)
y_test = np.random.rand(20, 1).astype(np.float32)

# 转换为 PyTorch 张量
X_test_tensor = torch.tensor(X_test)
y_test_tensor = torch.tensor(y_test)

# 评估模型
resnet.eval()  # 设置模型为评估模式
with torch.no_grad():
    test_outputs = resnet(X_test_tensor)
    test_loss = criterion(test_outputs, y_test_tensor)
    print(f'Test Loss: {test_loss.item():.4f}')

'''
注意事项
数据预处理：在实际应用中，图像数据通常需要经过预处理（如归一化）。
冻结/解冻层：如果你使用预训练的ResNet，可以选择冻结部分或全部层的参数，或者解冻一些层进行微调。
学习率和优化器：学习率和优化器的选择对训练结果有很大影响，可以根据需要进行调整。
通过上述步骤，你可以使用ResNet来训练一个非线性回归模型。希望这对你有所帮助！
'''